{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "440851a0-170d-4226-9857-f39f05cc6c70",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: http://mirrors.tencentyun.com/pypi/simple\n",
      "Requirement already satisfied: airllm in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (0.9.1)\n",
      "Collecting airllm\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/b5/36/d1cefb0725097e7ddf907783f31e9e17b191009978839a3d06598e72c41d/airllm-2.6-py3-none-any.whl (33 kB)\n",
      "Requirement already satisfied: transformers in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (4.35.0)\n",
      "Collecting transformers\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/20/0a/739426a81f7635b422fbe6cb8d1d99d1235579a6ac8024c13d743efa6847/transformers-4.36.2-py3-none-any.whl (8.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.2/8.2 MB\u001b[0m \u001b[31m1.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m0m\n",
      "\u001b[?25hRequirement already satisfied: tqdm in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from airllm) (4.66.1)\n",
      "Requirement already satisfied: torch in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from airllm) (2.1.0)\n",
      "Requirement already satisfied: accelerate in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from airllm) (0.24.1)\n",
      "Requirement already satisfied: safetensors in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from airllm) (0.4.0)\n",
      "Requirement already satisfied: optimum in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from airllm) (1.14.0)\n",
      "Requirement already satisfied: huggingface-hub in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from airllm) (0.17.3)\n",
      "Collecting scipy (from airllm)\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/69/f0/fb07a9548e48b687b8bf2fa81d71aba9cfc548d365046ca1c791e24db99d/scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (34.5 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m34.5/34.5 MB\u001b[0m \u001b[31m10.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: filelock in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers) (3.9.0)\n",
      "Collecting huggingface-hub (from airllm)\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/a0/0a/02ac0ae1047d97769003ff4fb8e6717024f3f174a5d13257415aa09e13d9/huggingface_hub-0.20.1-py3-none-any.whl (330 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m330.1/330.1 kB\u001b[0m \u001b[31m1.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: numpy>=1.17 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers) (1.24.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers) (23.2)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers) (6.0.1)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers) (2023.10.3)\n",
      "Requirement already satisfied: requests in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers) (2.31.0)\n",
      "Requirement already satisfied: tokenizers<0.19,>=0.14 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers) (0.14.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from huggingface-hub->airllm) (2023.10.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from huggingface-hub->airllm) (4.7.1)\n",
      "INFO: pip is looking at multiple versions of tokenizers to determine which version is compatible with other requirements. This could take a while.\n",
      "Collecting tokenizers<0.19,>=0.14 (from transformers)\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/ad/75/56230c5c65b226e707e1adbc759c19fdf1b20bb02c0276796b132c97118a/tokenizers-0.15.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m0m\n",
      "\u001b[?25hRequirement already satisfied: psutil in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from accelerate->airllm) (5.9.6)\n",
      "Requirement already satisfied: sympy in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from torch->airllm) (1.11.1)\n",
      "Requirement already satisfied: networkx in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from torch->airllm) (3.1)\n",
      "Requirement already satisfied: jinja2 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from torch->airllm) (3.1.2)\n",
      "Requirement already satisfied: coloredlogs in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from optimum->airllm) (15.0.1)\n",
      "Requirement already satisfied: datasets in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from optimum->airllm) (2.14.6)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from requests->transformers) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from requests->transformers) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from requests->transformers) (1.26.18)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from requests->transformers) (2023.7.22)\n",
      "Requirement already satisfied: sentencepiece!=0.1.92,>=0.1.91 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers[sentencepiece]>=4.26.0->optimum->airllm) (0.1.99)\n",
      "Requirement already satisfied: protobuf in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from transformers[sentencepiece]>=4.26.0->optimum->airllm) (4.25.0)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from coloredlogs->optimum->airllm) (10.0)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from datasets->optimum->airllm) (14.0.0)\n",
      "Requirement already satisfied: dill<0.3.8,>=0.3.0 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from datasets->optimum->airllm) (0.3.7)\n",
      "Requirement already satisfied: pandas in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from datasets->optimum->airllm) (2.0.3)\n",
      "Requirement already satisfied: xxhash in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from datasets->optimum->airllm) (3.4.1)\n",
      "Requirement already satisfied: multiprocess in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from datasets->optimum->airllm) (0.70.15)\n",
      "Requirement already satisfied: aiohttp in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from datasets->optimum->airllm) (3.8.6)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from jinja2->torch->airllm) (2.1.1)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from sympy->torch->airllm) (1.3.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from aiohttp->datasets->optimum->airllm) (23.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from aiohttp->datasets->optimum->airllm) (6.0.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from aiohttp->datasets->optimum->airllm) (4.0.3)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from aiohttp->datasets->optimum->airllm) (1.9.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from aiohttp->datasets->optimum->airllm) (1.4.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from aiohttp->datasets->optimum->airllm) (1.3.1)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from pandas->datasets->optimum->airllm) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from pandas->datasets->optimum->airllm) (2023.3.post1)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from pandas->datasets->optimum->airllm) (2023.3)\n",
      "Requirement already satisfied: six>=1.5 in /home/ubuntu/miniconda3/envs/ghostaienv/lib/python3.8/site-packages (from python-dateutil>=2.8.2->pandas->datasets->optimum->airllm) (1.16.0)\n",
      "Installing collected packages: scipy, huggingface-hub, tokenizers, transformers, airllm\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.17.3\n",
      "    Uninstalling huggingface-hub-0.17.3:\n",
      "      Successfully uninstalled huggingface-hub-0.17.3\n",
      "  Attempting uninstall: tokenizers\n",
      "    Found existing installation: tokenizers 0.14.1\n",
      "    Uninstalling tokenizers-0.14.1:\n",
      "      Successfully uninstalled tokenizers-0.14.1\n",
      "  Attempting uninstall: transformers\n",
      "    Found existing installation: transformers 4.35.0\n",
      "    Uninstalling transformers-4.35.0:\n",
      "      Successfully uninstalled transformers-4.35.0\n",
      "  Attempting uninstall: airllm\n",
      "    Found existing installation: airllm 0.9.1\n",
      "    Uninstalling airllm-0.9.1:\n",
      "      Successfully uninstalled airllm-0.9.1\n",
      "Successfully installed airllm-2.6 huggingface-hub-0.20.1 scipy-1.10.1 tokenizers-0.15.0 transformers-4.36.2\n"
     ]
    }
   ],
   "source": [
    "!pip install -U airllm transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "19eb8fee-ab17-4a54-9af2-ca809bd096b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>>> cache_utils installed\n",
      "saved layers already found in /home/ubuntu/.cache/huggingface/hub/models--mistralai--Mixtral-8x7B-v0.1/snapshots/58301445dc1378584211722b7ebf8743ec4e192b/splitted_model\n",
      "either BetterTransformer or attn_implementation='sdpa' is available, creating model directly\n",
      "either BetterTransformer or attn_implementation='sdpa' is available, creating model directly\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "cuda:0: 100%|██████████| 35/35 [04:29<00:00,  7.69s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "either BetterTransformer or attn_implementation='sdpa' is available, creating model directly\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "cuda:0: 100%|██████████| 35/35 [04:30<00:00,  7.73s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "either BetterTransformer or attn_implementation='sdpa' is available, creating model directly\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "cuda:0: 100%|██████████| 35/35 [04:29<00:00,  7.70s/it]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'<s> I like to think of'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from airllm import AutoModel\n",
    "\n",
    "MAX_LENGTH = 128\n",
    "# could use hugging face model repo id:\n",
    "model = AutoModel.from_pretrained(\"mistralai/Mixtral-8x7B-v0.1\")\n",
    "\n",
    "input_text = [\n",
    "        'I like',\n",
    "    ]\n",
    "\n",
    "input_tokens = model.tokenizer(input_text,\n",
    "    return_tensors=\"pt\",\n",
    "    return_attention_mask=False,\n",
    "    truncation=True,\n",
    "    max_length=MAX_LENGTH,\n",
    "    #padding=True\n",
    "    )\n",
    "\n",
    "generation_output = model.generate(\n",
    "    input_tokens['input_ids'].cuda(),\n",
    "    max_new_tokens=3,\n",
    "    use_cache=True,\n",
    "    return_dict_in_generate=True)\n",
    "\n",
    "model.tokenizer.decode(generation_output.sequences[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db4d8871-7e30-4eb8-b2f9-0310409c71d7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
